import pandas as pd
from sklearn.model_selection import KFold
from Ensemble import Ensemble, MultiClassifier
import re as re
import numpy as np

trainFile = "train.csv"
testFile = "test.csv"
Ids = object()

def DataClean(ds,saveID=False):
    global Ids
    dropFeat = ['PassengerId','Name','Ticket','Cabin','SibSp','Parch']

    ds['Age'] = ds['Age'].fillna(ds['Age'].mean())
    ds['Embarked'] = ds['Embarked'].fillna('S')
    ds['Fare'] = ds['Fare'].fillna(ds['Fare'].mean())

    ds['Sex'] = ds['Sex'].map({'male':0,'female':1}).astype(int)
    ds['Embarked'] = ds['Embarked'].map({'Q':0,'S':1,'C':2}).astype(int)

    ds['Title'] = ds['Name'].apply(lambda x: re.search(' ([A-Za-z]+)\.',x).group(1))
    ds['Title'].replace(['Capt', 'Col', 'Countess', 'Don', 'Dr', 'Jonkheer', 'Lady', 'Major', 'Rev', 'Sir', 'Dona'],'Rare')
    ds['Title'].map({'Mlle':'Miss', 'Ms':'Miss', 'Mme':'Mrs'})
    ds['Title'] = ds['Title'].map({'Mr':1,'Miss':2,'Mrs':3,'Master':4,'Rare':5}).fillna(0)

    ds['FSize'] = ds['SibSp'] + ds['Parch'] + 1
    ds['IsAlone'] = 0
    ds.loc[ds['FSize'] == 1,'IsAlone'] = 1

    #todo : what if the distribution of train and test dataset is different??
    ds.loc[ds['Fare'] <= 7.91, 'Fare'] = 0
    ds.loc[(ds['Fare'] > 7.91) & (ds['Fare'] <= 14.454), 'Fare'] = 1
    ds.loc[(ds['Fare'] > 14.454) & (ds['Fare'] <= 31), 'Fare'] = 2
    ds.loc[ds['Fare'] > 31, 'Fare'] = 3
    ds['Fare'] = ds['Fare'].astype(int)
    
    # Mapping Age
    ds.loc[ds['Age'] <= 16, 'Age'] = 0
    ds.loc[(ds['Age'] > 16) & (ds['Age'] <= 32), 'Age'] = 1
    ds.loc[(ds['Age'] > 32) & (ds['Age'] <= 48), 'Age'] = 2
    ds.loc[(ds['Age'] > 48) & (ds['Age'] <= 64), 'Age'] = 3
    ds.loc[ds['Age'] > 64, 'Age'] = 4


    if saveID:
        Ids = ds['PassengerId']
    ds = ds.drop(dropFeat,1)
    return ds

def ReadData():
    data_train = DataClean(pd.read_csv(trainFile))
    data_test = DataClean(pd.read_csv(testFile),saveID = True)
    return data_train, data_test

def SelectModel(data_train,data_test):
    acc = 0
    NC = 5
    kf = KFold(NC, random_state = 0)
    X = data_train.iloc[0:,1:]
    y = data_train.iloc[0:,0]
    m = MultiClassifier()
    ens = Ensemble()
    #oof_train = np.zeros((data_train.shape[0],len(m.models)))
    #oof_test = np.zeros((data_test.shape[0],))
    #oof_test_skf = np.empty((NC,data_test.shape[0]))

    for train_index,test_index in kf.split(X):
        X_train,X_test = X.iloc[train_index], X.iloc[test_index]
        y_train,y_test = y.iloc[train_index], y.iloc[test_index]

        m.train(X_train,y_train)
        new_train = m.predict(X_test)
        ens.train(new_train,y_test)




    print(acc / NC)
    return (m,ens)

def GenResult(m,data_test):
    global Ids
    pred = m[1].predict(m[0].predict(data_test))
    Prediction = pd.DataFrame({'PassengerId': Ids,'Survived':pred})
    Prediction.to_csv('prediction.csv', index=False)

if __name__ == "__main__":
    data_train,data_test = ReadData()
    mdls = SelectModel(data_train,data_test)
    GenResult(mdls,data_test)

